'''
Created on Aug 31, 2009

@author: mkiyer
'''

import logging
import sys

hg18_chrom_lengths = {'chr1': 247249719,
                      'chr10': 135374737,
                      'chr11': 134452384,
                      'chr12': 132349534,
                      'chr13': 114142980,
                      'chr14': 106368585,
                      'chr15': 100338915,
                      'chr16': 88827254,
                      'chr17': 78774742,
                      'chr18': 76117153,
                      'chr19': 63811651,
                      'chr2': 242951149,
                      'chr20': 62435964,
                      'chr21': 46944323,
                      'chr22': 49691432,
                      'chr3': 199501827,
                      'chr4': 191273063,
                      'chr5': 180857866,
                      'chr6': 170899992,
                      'chr7': 158821424,
                      'chr8': 146274826,
                      'chr9': 140273252,
                      'chrX': 154913754,
                      'chrY': 57772954
                      }

class Gene(object):    
    def __init__(self):
        self.exons = []
        self.introns = []
        self.utr5 = None
        self.utr3 = None
        self.aliases = []
        
def parse_ucsc_genes(fhd):
    '''
    name varchar(255) NOT NULL default '',
    chrom varchar(255) NOT NULL default '',
    strand char(1) NOT NULL default '',
    txStart int(10) unsigned NOT NULL default '0',
    txEnd int(10) unsigned NOT NULL default '0',
    cdsStart int(10) unsigned NOT NULL default '0',
    cdsEnd int(10) unsigned NOT NULL default '0',
    exonCount int(10) unsigned NOT NULL default '0',
    exonStarts longblob NOT NULL,
    exonEnds longblob NOT NULL,
    proteinID varchar(40) NOT NULL default '',
    alignID varchar(255) NOT NULL default '',
    '''
    for line in fhd:
        fields = line.strip().split('\t')
        g = Gene()
        g.name = fields[0]
        g.chrom = fields[1]
        if g.chrom not in hg18_chrom_lengths:
            continue
        g.strand = fields[2]
        g.txstart = int(fields[3])
        g.txend = int(fields[4])
        g.cdsstart = int(fields[5])
        g.cdsend = int(fields[6])
        g.exon_count = int(fields[7])
        g.exon_starts = map(int, fields[8].split(',')[:-1])
        g.exon_ends = map(int, fields[9].split(',')[:-1])
        g.exons = zip(g.exon_starts, g.exon_ends)
        g.introns = zip(g.exon_ends, g.exon_starts[1:])
        # figure out the UTRs
        if g.txstart != g.cdsstart:
            if g.strand == '+':
                g.utr5 = (g.txstart, g.cdsstart)
            elif g.strand == '-':
                g.utr3 = (g.txstart, g.cdsstart)
            else:
                logging.critical("Gene strand is not '+' or '-': %s" % g.name)
                sys.exit(1)
        if g.txend != g.cdsend:
            if g.strand == '+':
                g.utr3 = (g.cdsend, g.txend)
            elif g.strand == '-':
                g.utr5 = (g.cdsend, g.txend)
            else:
                logging.critical("Gene strand is not '+' or '-': %s" % g.name)
                sys.exit(1)
        yield g

def plot_hists(x1, y1, x2, y2, outfile, x_name="group1", y_name="group2"):
    import pylab as P
    # histogram
    P.figure()

    nbins = 100
    #bins = np.arange(0, 1000, 10)
    
    P.subplot(121)
    # the histogram of the data
    #n, bins, patches = P.hist([x,y], nbins, alpha=1)
    n1, bins1, patches1 = P.hist(x1, bins=nbins, facecolor='g', alpha=0.5)        
    n2, bins2, patches2 = P.hist(y1, bins=bins1, facecolor='r', alpha=0.5)
    P.legend([patches1[0], patches2[0]], [x_name, y_name])
    P.grid()
    P.xlabel('normalized coverage')
    P.ylabel('bin count')
    P.title("First PolyA site")        
    P.subplot(122)
    n1, bins1, patches1 = P.hist(x2, bins=nbins, facecolor='g', alpha=0.5)        
    n2, bins2, patches2 = P.hist(y2, bins=bins1, facecolor='r', alpha=0.5)
    P.legend([patches1[0], patches2[0]], [x_name, y_name])
    P.grid()
    P.xlabel('normalized coverage')
    P.ylabel('bin count')
    P.title("Rest of 3'UTR")
    P.savefig(outfile)
    P.close()


if __name__ == '__main__':
    # read ucsc genes files
    logging.debug("loading ucsc known genes")
    genes = [x for x in parse_ucsc_genes(open('knownGene.txt'))]

    exonsizes = []
    intronsizes = []
    
    for g in genes:
        if g.cdsstart == g.cdsend:
            continue
        for e in g.exons:
            exonsizes.append(e[1]-e[0])
        for i in g.introns:
            intronsizes.append(i[1] - i[0])
            
    
    import pylab as P
    bins = range(0, 3000, 10)
    P.figure()
    n1, bins1, patches1 = P.hist(exonsizes, bins=bins, facecolor='g')
    P.xlim(xmax=2000)
    P.savefig('exonprofile.png')
    P.close()

    bins = range(0, 5000, 10)
    P.figure()
    n1, bins1, patches1 = P.hist(intronsizes, bins=bins, facecolor='g')
    P.xlim(xmax=2000)
    P.savefig('intronprofile.png')
    P.close()

